import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import math
import matplotlib.pyplot as plt
import numpy as np
import copy



# *******************文本嵌入层*******************


# embedding = nn.Embedding(10, 3)
# input1 = torch.LongTensor([[1, 2, 4, 5], [4, 3, 2, 9]])
# print(embedding(input1))


# embedding = nn.Embedding(10, 3, padding_idx=0)
# input1 = torch.LongTensor([0, 2, 0, 5])
# print(embedding(input1))


# 构建Embedding类来实现文本嵌入层
class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        # d_model：词嵌入的维度
        # vocab：词表的大小
        super(Embeddings, self).__init__()
        # 定义Embedding层
        self.lut = nn.Embedding(vocab, d_model)
        # 将参数传入类中
        self.d_model = d_model

    def forward(self, x):
        # x：代表输入进模型的文本通过词汇映射后的数字张量
        # self.lut(x)：将数字张量通过Embedding层映射成词嵌入表示
        # *math.sqrt(self.d_model)：对词嵌入表示进行缩放以保持数值稳定性。
        return self.lut(x) * math.sqrt(self.d_model)


d_model = 512
vocab = 1000

x = Variable(torch.LongTensor([[100, 2, 421, 508], [491, 998, 1, 221]]))
emb = Embeddings(d_model, vocab)
embr = emb(x)
# print("embr:", embr)
# print("embr.size():", embr.size())



# *******************位置编码器类*******************

# m = nn.Dropout(p=0.2)
# input1 = torch.randn(4, 5)
# output = m(input1)
# print(output)

# x = torch.tensor([1, 2, 3, 4])
# print(x)
# print(x.size())
# y = torch.unsqueeze(x, 0)
# print(y)
# print(y.size())
# z = torch.unsqueeze(x, 1)
# print(z)
# print(z.size())


# 构建位置编码器的类
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout, max_len=5000):
        # d_model：代表词嵌入的维度
        # dropout：代表Dropout层的置零比率
        # max_len：代表每段句子的最大长度
        super(PositionalEncoding, self).__init__()

        # 实例化Dropout层，使用nn.Dropout层来防止过拟合
        self.dropout = nn.Dropout(p=dropout)

        # 初始化一个位置编码矩阵，大小是max_len * d_model
        pe = torch.zeros(max_len, d_model)

        # 初始化一个绝对位置矩阵，torch.arange(0, max_len)会生成一个从0到max_len-1的连续整数序列，
        # 这个序列的形状是(max_len,)。然后，unsqueeze(1)会将这个一维张量转换为一个二维张量，形状变为(max_len, 1)。
        position = torch.arange(0, max_len).unsqueeze(1)

        # 定义一个变化矩阵div_term，跳跃式的初始化
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))

        # 将绝对位置矩阵position与变化矩阵div_term相乘，得到一个矩阵，
        # 这个矩阵的每一行对应一个位置，每一列对应一个维度，每一行的值是position * div_term。
        # 最后，将这个矩阵的奇数列的值用torch.sin函数计算，偶数列的值用torch.cos函数计算，得到位置编码矩阵pe。
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        # 位置编码矩阵pe的形状是(max_len, d_model)，其中max_len是每段句子的最大长度，d_model是词嵌入的维度。
        # 将二维张量扩充成三维张量，unsqueeze(0)会将pe的形状从(max_len, d_model)变为(1, max_len, d_model)。
        pe = pe.unsqueeze(0)

        # 将位置编码矩阵注册成模型的buffer，这个buffer不是模型中的参数，不跟随优化器更新，用来保存一些固定的数据
        # 注册成buffer后我们就可以在模型保存后重新加载的时候，将这个位置编码器和模型参数一起加载出来
        self.register_buffer('pe', pe)

    def forward(self, x):
        # 输入的x代表词嵌入矩阵
        # 将词嵌入矩阵和位置编码矩阵相加，然后通过一个Dropout层
        # Variable被用来创建一个位置编码矩阵的Variable对象，这个对象会被添加到词嵌入矩阵x上。
        # self.pe[:, :x.size(1)]会取出位置编码矩阵pe的所有行和从第0列开始，到x的第二个维度大小结束的列，也就是与x的第二个维度大小相同的部分。
        # requires_grad=False表示这个Variable对象不需要计算梯度，也就是说，在反向传播过程中，这个Variable对象的梯度会被忽略。
        # Variable对象被创建后，会被添加到词嵌入矩阵x上，然后通过一个Dropout层。这个操作是为了防止过拟合，
        # Dropout层会随机地将输入的一部分元素置零。最后，self.dropout(x)会返回一个经过Dropout层处理后的词嵌入矩阵。
        x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)
        return self.dropout(x)


d_model = 512
dropout = 0.1
max_len = 60

x = embr
pe = PositionalEncoding(d_model, dropout, max_len)
pe_result = pe(x)
# print(pe_result)
# print(pe_result.shape)


# 第一步设置一个画布
plt.figure(figsize=(15, 5))

# 实例化PositionalEncoding类对象，词嵌入维度给20，置零比率设置为0
pe = PositionalEncoding(20, 0)

# 向pe中传入一个全零初始化的x，相当于展示pe
y = pe(Variable(torch.zeros(1, 100, 20)))

plt.plot(np.arange(100), y[0, :, 4:8].data.numpy())
plt.legend(["dim %d" % p for p in [4, 5, 6, 7]])
# plt.show()



# print(np.triu(np.ones((3, 3)), k=1))
# print(np.triu([[1,2,3],[4,5,6],[7,8,9]], k=1))

def subsequent_mask(size):
    """生成向后遮掩的掩码张量，参数size是掩码张量最后两个维度的大小，它的最后两维形成一个方阵"""
    # 在函数中，首先定义掩码张量的形状
    attn_shape = (1, size, size)

    # 然后使用np.ones方法向这个形状中添加1元素，形成上三角阵，最后为了节约空间，
    # 再使其中的数据类型变为无符号8位整形unit8
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')

    # 最后将numpy类型转化为torch中的tensor，内部做一个1 - 的操作，
    # 在这个其实是做了一个三角阵的反转，subsequent_mask中的每个元素都会被1减，
    # 如果是0，subsequent_mask中的该位置由0变成1
    # 如果是1，subsequent_mask中的该位置由1变成0
    return torch.from_numpy(1 - subsequent_mask)


size = 5
mask = subsequent_mask(size)
# print(mask)




# x = Variable(torch.randn(5, 5))
# print(x)

# mask = Variable(torch.zeros(5, 5))
# print(mask)

# y = x.masked_fill(mask == 0, -1e9)
# print(y)


def attention(query, key, value, mask=None, dropout=None):
    # query, key, value：代表注意力的三个输入张量
    # mask：掩码张量
    # dropout：传入的Dropout实例化对象
    # 首先将query的最后一个维度提取出来，代表的是词嵌入的维度
    d_k = query.size(-1)

    # 按照注意力计算公式，将query和key的转置进行矩阵乘法，然后除以缩放稀疏
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

    # 判断是否使用掩码张量
    if mask is not None:
        # 利用masked_fill方法，将掩码张量和0进行位置的意义比较，如果等于0，替换成一个非常小的数
        scores = scores.masked_fill(mask == 0, -1e9)

    # 对scores的最后一个维度上进行softmax操作
    p_attn = F.softmax(scores, dim=-1)

    # 判断是否使用dropout
    if dropout is not None:
        p_attn = dropout(p_attn)

    # 最后一步完成p_attn和value张量的乘法，并返回query注意力表示
    return torch.matmul(p_attn, value), p_attn

query = key = value = pe_result
mask = Variable(torch.zeros(2, 4, 4))
attn, p_attn = attention(query, key, value, mask=mask)
# print('attn:', attn)
# print(attn.shape)
# print('p_attn:', p_attn)
# print(p_attn.shape)


# 需要使用clone函数将它们一同初始化到一个网络层列表对象中
def clones(module, N):
    # module：代表要克隆的目标网络层
    # N：将module克隆几个
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


# 实现多头注意力机制的类
class MultiHeadedAttention(nn.Module):
    def __init__(self, head, embedding_dim, dropout=0.1):
        # head：代表几个头的参数
        # embedding_dim：代表词嵌入的维度
        # dropout：进行Dropout操作时，置零的比率
        super(MultiHeadedAttention, self).__init__()

        # 要确认一个事实：多头的数量head需要整除词嵌入的维度embedding_dim
        assert embedding_dim % head == 0

        # 得到每个头获得的词向量的维度
        self.d_k = embedding_dim // head

        self.head = head
        self.embedding_dim = embedding_dim

        # 获得线性层，要获得4个，分别是Q,K,V以及最终的输出线性层
        self.linears = clones(nn.Linear(embedding_dim, embedding_dim), 4)

        # 初始化注意力张量
        self.attn = None

        # 初始化dropout对象
        self.dropout = nn.Dropout(p=dropout)


    def forward(self, query, key, value, mask=None):
        # query, key, value是注意力机制的三个输入张量, mask代表掩码张量
        # 首先判断是否使用掩码张量
        if mask is not None:
            # 使用squeeze将掩码张量进行维度扩充，代表多头中的第n个头
            mask = mask.unsqueeze(1)

        # 得到batch_size
        batch_size = query.size(0)

        # 首先使用zip将网络层和输入数据连接在一起，模型的输出利用view和transpose进行维度和形状的
        query, key, value = \
            [model(x).view(batch_size, -1, self.head, self.d_k).transpose(1, 2)
            for model, x in zip(self.linears, (query, key, value))]

        # 将每个头的输出传入到注意力层
        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)

        # 得到每个头的计算结果是4维张量， 需要进行形状的转换
        # 前面已经将1,2两个维度进行过转置，在这里要重新转置回来
        # 注意: 经历了transpose()方法后，必须要使用contiguous方法，不然无法使用view()方法
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.head * self.d_k)

        # 最后将x输入线性层列表中的最后一个线性层中进行处理，得到最终的多头注意力结构输出
        return self.linears[-1](x)


# 实例化若干参数
head = 8
embedding_dim = 512
dropout = 0.2

# 若干输入参数的初始化
query = key = value = pe_result
mask = Variable(torch.zeros(2, 4, 4))

mha = MultiHeadedAttention(head, embedding_dim, dropout)
mha_result = mha(query, key, value, mask)
# print(mha_result)
# print(mha_result.shape)


# 通过类PositionwiseFeedForward来实现前馈全连接层
class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        """初始化函数有三个输入参数分别是d_model, d_ff, 和dropout=0.1，第一个是线性层的输入维
        因为我们希望输入通过前馈全连接层后输入和输出的维度不变。第二个参数d_ff就是第二个线性层
        最后一个是dropout置0比率。"""
        super(PositionwiseFeedForward, self).__init__()

        # 首先按照我们预期使用nn实例化了两个线性层对象，self.w1和self.w2
        # 它们的参数分别是d_model, d_ff和d_ff, d_model
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)

        # 然后使用nn的Dropout实例化了对象self.dropout
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        """输入参数为x，代表来自上一层的输出"""
        # 首先经过第一个线性层，然后使用Functional中relu函数进行激活，
        # 之后再使用dropout进行随机置0，最后通过第二个线性层w2，返回最终结果。
        return self.w2(self.dropout(F.relu(self.w1(x))))
    

d_model = 512
d_ff = 64
dropout = 0.2

x = mha_result
ff = PositionwiseFeedForward(d_model, d_ff, dropout)
ff_result = ff(x)
# print(ff_result)
# print(ff_result.shape)




# 构建规范化层的类
class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        # features：代表词嵌入的维度
        # eps：一个足够小的正数，用来在规范化计算公式的分母中，防止除零操作
        super(LayerNorm, self).__init__()
        # 初始化两个参数张量a2, b2，用于对结果做规范化操作计算
        # 将其用nn.Parameter进行封装，代表它们也是模型中的参数
        self.a2 = nn.Parameter(torch.ones(features))
        self.b2 = nn.Parameter(torch.zeros(features))
        self.eps = eps
    def forward(self, x):
        # x：代表上一层网络的输出
        # 首先对x进行最后一个维度上的求均值操作，同时保持输出维度和输入维度一致
        mean = x.mean(-1, keepdim=True)
        # 接着对x进行最后一个维度上的求标准差的操作，同时保持输出维度和输入维度一致
        std = x.std(-1, keepdim=True)
        # 按照规范化公式进行计算并返回
        return self.a2 * (x - mean) / (std + self.eps) + self.b2

features = d_model = 512
eps = 1e-6

x = ff_result
ln = LayerNorm(features, eps)
ln_result = ln(x)
# print(ln_result)
# print(ln_result.shape)


# 构建子层连接结构的类
class SublayerConnection(nn.Module):
    def __init__(self, size, dropout=0.1):
        # size：代表词嵌入的维度
        # dropout：进行Dropout操作的置零比率
        super(SublayerConnection, self).__init__()
        # 实例化一个规范化层的对象
        self.norm = LayerNorm(size)
        # 实例化一个Dropout对象
        self.dropout = nn.Dropout(p=dropout)
        self.size = size
    def forward(self, x, sublayer):
        # x：代表上一层传入的张量
        # sublayer：该子层连接中子层函数
        # 首先将x进行规范化，然后送入子层函数中处理，处理结果进入dropout层，最后进行残差连接
        return x + self.dropout(sublayer(self.norm(x)))
    
size = d_model = 512
head = 8
dropout = 0.2
x = pe_result
mask = torch.zeros(2, 4, 4)
self_attn = MultiHeadedAttention(head, d_model)
sublayer = lambda x: self_attn(x, x, x, mask)
sc = SublayerConnection(size, dropout)
sc_result = sc(x, sublayer)
# print(sc_result)
# print(sc_result.shape)



# 使用EncoderLayer类实现编码器层
class EncoderLayer(nn.Module):
    def __init__(self, size, self_attn, feed_forward, dropout):
        """它的初始化函数参数有四个，分别是size，其实就是我们词嵌入维度的大小，它也将作为我们编码器层的输入输出维度大小。
        第二个是self_attn，之后我们将传入多头自注意力子层实例化对象，并且是自注意力机制，
        第三个是feed_froward，之后我们将传入前馈全连接层实例化对象，最后一个是置0比率dropout"""
        super(EncoderLayer, self).__init__()
        # 首先将self_attn和feed_forward传入其中.
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        # 如图所示，编码器层中有两个子层连接结构，所以使用clones函数进行克隆
        self.sublayer = clones(SublayerConnection(size, dropout), 2)
        # 把size传入其中
        self.size = size
    def forward(self, x, mask):
        """forward函数中有两个输入参数，x和mask，分别代表上一层的输出，和掩码张量mask。"""
        # 里面就是按照结构图左侧的流程。首先通过第一个子层连接结构，其中包含多头自注意力子层，
        # 然后通过第二个子层连接结构，其中包含前馈全连接子层。最后返回结果。
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)

size = 512
head = 8
d_model = 512
d_ff = 64
x = pe_result
dropout = 0.2
self_attn = MultiHeadedAttention(head, d_model)
ff = PositionwiseFeedForward(d_model, d_ff, dropout)
mask = Variable(torch.zeros(2, 4, 4))

el = EncoderLayer(size, self_attn, ff, dropout)
el_result = el(x, mask)
# print(el_result)
# print(el_result.shape)


# 使用Encoder类来实现编码器
class Encoder(nn.Module):
    def __init__(self, layer, N):
        """初始化函数的两个参数分别代表编码器层和编码器层的个数"""
        super(Encoder, self).__init__()
        # 首先使用clones函数克隆N个编码器层放在self.layers中
        self.layers = clones(layer, N)
        # 再初始化一个规范化层，它将用在编码器的最后面.
        self.norm = LayerNorm(layer.size)
    def forward(self, x, mask):
        """forward函数的输入和编码器层相同，x代表上一层的输出，mask代表掩码张量"""
        # 首先就是对我们克隆的编码器层进行循环，每次都会得到一个新的x，
        # 这个循环的过程，就相当于输出的x经过了N个编码器层的处理.
        # 最后再通过规范化层的对象self.norm进行处理，最后返回结果.
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)
    
# 第一个实例化参数layer，它是一个编码器层的实例化对象，因此需要传入编码器层的参数
# 又因为编码器层中的子层是不共享的，因此需要使用深度拷贝各个对象.
size = 512
head = 8
d_model = 512
d_ff = 64
c = copy.deepcopy
attn = MultiHeadedAttention(head, d_model)
ff = PositionwiseFeedForward(d_model, d_ff, 0.2)
dropout = 0.2
layer = EncoderLayer(size, c(attn), c(ff), dropout)
# 编码器中编码器层的个数N
N = 8
mask = Variable(torch.zeros(2, 4, 4))


en = Encoder(layer, N)
en_result = en(x, mask)
# print(en_result)
# print(en_result.shape)




# 使用DecoderLayer的类实现解码器层
class DecoderLayer(nn.Module):
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        """初始化函数有5个，分别是size，代表词嵌入的维度大小，同时也代表解码器层的尺寸，    
        第二个是self_attn，多头自注意力对象，也就是说这个注意力机制需要Q=K=V，  
        第三个是src_attn，多头注意力对象，这里Q!=K=V，第四个是前馈全连接层对象，最后就是dropout置0比率"""          
        super(DecoderLayer, self).__init__()

        # 在初始化函数中，主要就是将这些传入到类中
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward

        # 按照结构图使用clones函数克隆三个子层连接对象.
        self.sublayer = clones(SublayerConnection(size, dropout), 3)

    def forward(self, x, memory, source_mask, target_mask):
        """forward函数中的参数有4个，分别是来自上一层的输入x，
        来自编码器层的语义存储变量memory，以及源数据掩码张量和目标数据掩码张量。"""
        # 将memory表示成m方便之后使用
        m = memory

        # 将x传入第一个子层结构，第一个子层结构的输入分别是x和self - attn，因为是自注意力机制，
        # 最后一个参数是目标数据掩码张量，这时要对目标数据进行遮掩，因为此时模型可能还没有生成任何数据，    
        # 比如在解码器准备生成第一个字符或词汇时，我们其实已经传入了第一个字符以便计算损失，
        # 但是我们不希望在生成第一个字符时模型能利用这个信息，因此我们会将其遮掩，同样生成第二个字符时，第一个字符的信息也会被遮掩，以此类推，直到生成完整的目标序列。
        # 模型只能使用第一个字符或词汇信息，第二个字符以及之后的信息都不允许被模型使用。
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, target_mask))

        # 接着进入第二个子层，这个子层中常规的注意力机制，q是输入x；k，v是编码层输出memory，
        # 同样也传入source_mask，但是进行源数据遮掩的原因并非是抑制信息泄漏，而是遮蔽掉对结果没有帮助的位置。   
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, source_mask))

        # 最后进入第三个子层，也就是前馈全连接子层
        return self.sublayer[2](x, self.feed_forward)


head = 8
size = 512
d_model = 512
d_ff = 64
dropout = 0.2
self_attn = src_attn = MultiHeadedAttention(head, d_model, dropout)

ff = PositionwiseFeedForward(d_model, d_ff, dropout)

# x是来自目标数据的词嵌入表示，但形式和源数据的词嵌入表示相同，这里使用pe_result充当.
x = pe_result
# memory是来自编码器的输出
memory = en_result
# 实际中source_mask和target_mask并不相同，这里为方便计算使他们都为mask
mask = Variable(torch.zeros(2, 4, 4))
source_mask = target_mask = mask

dl = DecoderLayer(size, self_attn, src_attn, ff, dropout)
dl_result = dl(x, memory, source_mask, target_mask)
# print(dl_result)
# print(dl_result.shape)



# 使用类Decoder来实现解码器
class Decoder(nn.Module):
    def __init__(self, layer, N):
        """初始化函数的参数有两个，第一个就是解码器层layer，第二个是解码器层的个数N。"""
        super(Decoder, self).__init__()

        # 首先使用clones方法克隆了N个layer，然后实例化了一个规范化层。
        # 因为数据走过了所有的解码器层后最后要做规范化处理。
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, memory, source_mask, target_mask):
        """forward函数中的参数有4个，x代表目标数据的嵌入表示，memory是编码器层的输出，
        source_mask，target_mask代表源数据和目标数据的掩码张量"""

        # 然后就是对每个层进行循环，当然这个循环就是变量x通过每一个层的处理，
        # 得出最后的结果，再进行一次规范化返回即可。
        for layer in self.layers:
            x = layer(x, memory, source_mask, target_mask)
        return self.norm(x)
    

size = 512
d_model = 512
head = 8
d_ff = 64
dropout = 0.2
c = copy.deepcopy
attn = MultiHeadedAttention(head, d_model)
ff = PositionwiseFeedForward(d_model, d_ff, dropout)
layer = DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout)
N = 8

# 输入参数与解码器层的输入参数相同
x = pe_result
memory = en_result
mask = Variable(torch.zeros(2, 4, 4))
source_mask = target_mask = mask

de = Decoder(layer, N)
de_result = de(x, memory, source_mask, target_mask)
# print(de_result)
# print(de_result.shape)




# 将线性层和softmax计算层一起实现，因为二者的共同目标是生成最后的结构
# 因此把类的名字叫做Generator，生成器类
class Generator(nn.Module):
    def __init__(self, d_model, vocab_size):
        """初始化函数的输入参数有两个，d_model代表词嵌入维度，vocab_size代表词表大小。"""
        super(Generator, self).__init__()
        # 首先就是使用nn中的预定义线性层进行实例化，得到一个对象self.project等待使用，
        # 这个线性层的参数有两个，就是初始化函数传进来的两个参数：d_model，vocab_size
        self.project = nn.Linear(d_model, vocab_size)

    def forward(self, x):
        """前向逻辑函数中输入是上一层的输出张量x"""
        # 在函数中，首先使用上一步得到的self.project对x进行线性变化，
        # 然后使用F中已经实现的log_softmax进行的softmax处理。
        # 在这里之所以使用log_softmax是因为和我们这个pytorch版本的损失函数实现有关，在其他版本
        # log_softmax就是对softmax的结果又取了对数，因为对数函数是单调递增函数，
        # 因此对最终我们取最大的概率值没有影响。最后返回结果即可。
        return F.log_softmax(self.project(x), dim=-1)
    
# 词嵌入维度是512维
d_model = 512
# 词表大小是1000
vocab_size = 1000

# 输入x是上一层网络的输出，我们使用来自解码器层的输出
x = de_result

# 调用
gen = Generator(d_model, vocab_size)
gen_result = gen(x)
# print(gen_result)
# print(gen_result.shape)




# 使用EncoderDecoder类来实现编码器 - 解码器结构
class EncoderDecoder(nn.Module):
    def __init__(self, encoder, decoder, source_embed, target_embed, generator):
        """初始化函数有5个参数，分别是编码器对象，解码器对象，
        源数据嵌入函数，目标数据嵌入函数，以及输出部分的类别生成器对象"""
        super(EncoderDecoder, self).__init__()
        # 将参数传入到类中
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = source_embed
        self.tgt_embed = target_embed
        self.generator = generator

    def forward(self, source, target, source_mask, target_mask):
        """在forward函数中，有四个参数，source代表源数据，target代表目标数据，
        source_mask和target_mask代表对应的掩码张量"""
        # 在函数中，将source，source_mask传入编码函数，得到结果后，
        # 与source_mask，target，和target_mask一同传给解码函数.
        return self.decode(self.encode(source, source_mask), source_mask, target, target_mask)

    def encode(self, source, source_mask):
        """编码函数，以source和source_mask为参数"""
        # 使用src_embed对source做处理，然后和source_mask一起传给self.encoder
        return self.encoder(self.src_embed(source), source_mask)

    def decode(self, memory, source_mask, target, target_mask):
        """解码函数，以memory即编码器的输出，source_mask，target，target_mask为参数"""
        # 使用tgt_embed对target做处理，然后和source_mask，target_mask，memory一起传给self.decoder
        return self.decoder(self.tgt_embed(target), memory, source_mask, target_mask)
    

vocab_size = 1000
d_model = 512
encoder = en
decoder = de
source_embed = nn.Embedding(vocab_size, d_model)
target_embed = nn.Embedding(vocab_size, d_model)
generator = gen

# 假设源数据与目标数据相同，实际中并不相同
source = target = Variable(torch.LongTensor([[100, 2, 421, 508], [491, 998, 1, 221]]))
# 假设src_mask与tgt_mask相同，实际中并不相同
source_mask = target_mask = Variable(torch.zeros(2, 4, 4))
# 调用
ed = EncoderDecoder(encoder, decoder, source_embed, target_embed, generator)
ed_result = ed(source, target, source_mask, target_mask)
# print(ed_result)
# print(ed_result.shape)




def make_model(source_vocab, target_vocab, N=6,
               d_model=512, d_ff=2048, head=8, dropout=0.1):
    """该函数用来构建模型，有7个参数，分别是源数据特征(词汇)总数，目标数据特征(词汇)总数，
    编码器和解码器堆叠数，词向量映射维度，前馈全连接网络中变换矩阵的维度，
    多头注意力结构中的头数，以及置零比率dropout。"""
    # 首先得到一个深度拷贝命令，接下来很多结构都需要进行深度拷贝，
    # 来保证他们彼此之间相互独立，不受干扰
    c = copy.deepcopy
    # 实例化了多头注意力类，得到对象attn
    attn = MultiHeadedAttention(head, d_model)
    # 然后实例化前馈全连接类，得到对象ff
    ff = PositionwiseFeedForward(d_model, d_ff, dropout)
    # 实例化位置编码类，得到对象position
    position = PositionalEncoding(d_model, dropout)
    # 根据结构图，最外层是EncoderDecoder，在EncoderDecoder中，
    # 分别是编码器层，解码器层，源数据Embedding层和位置编码组成的有序结构，
    # 目标数据Embedding层和位置编码组成的有序结构，以及类别生成器层。
    # 在编码器层中有attenton子层以及前馈全连接层。
    # 在解码器层中有两个attention子层以及前馈全连接层。
    model = EncoderDecoder(
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
        nn.Sequential(Embeddings(d_model, source_vocab), c(position)),
        nn.Sequential(Embeddings(d_model, target_vocab), c(position)),
        Generator(d_model, target_vocab))
    
    # 模型结构完成后，接下来就是初始化模型中的参数，比如线性层中的变换矩阵
    # 这里一但判断参数的维度大于1，则会将其初始化成一个服从均匀分布的矩阵，
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    return model


source_vocab = 11
target_vocab = 11
N = 6
# 其他参数都使用默认值

if __name__ == '__main__':
    res = make_model(source_vocab, target_vocab, N)
    print(res)

